package com.immortal.ai.controller;

import com.immortal.ai.etity.vo.MessageVO;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.ArrayList;
import java.util.List;

/**
 * @author ywh
 * @version 1.0
 * @description: TODO
 * @date 2025/6/9 17:46
 */
@RequiredArgsConstructor
@RequestMapping("/ai")
@RestController
public class ChatController {

    public final ChatClient chatClient;

    public final ChatClient fortuneTellingChatClient;

    public final ChatClient functionCallChatClient;

    public final ChatClient pdfChatClient;

    public final ChatMemory chatMemory;

    public final VectorStore vectorStore;

    @GetMapping(value = "/chat",produces = "text/html;charset=utf-8")
    public String chat(@RequestParam("prompt") String prompt) {
        return chatClient.prompt()
                .user(prompt)
                .call()
                .content();
    }


    @GetMapping(value = "/promptChat",produces = "text/html;charset=utf-8")
    public String promptChat(@RequestParam("prompt") String prompt,@RequestParam("chatId") String chatId) {
        return fortuneTellingChatClient.prompt()
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,chatId))
                .call()
                .content();
    }


    @GetMapping(value = "/funPromptChat",produces = "text/html;charset=utf-8")
    public String functionChat(@RequestParam("prompt") String prompt,@RequestParam("chatId") String chatId) {
        return functionCallChatClient.prompt()
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,chatId))
                .call()
                .content();
    }

    @GetMapping(value = "/history")
    public List<MessageVO> getChatHistory(@RequestParam("chatId") String chatId) {
        List<MessageVO> messageVOList = new ArrayList<>();
        chatMemory.get(chatId,Integer.MAX_VALUE).forEach(message -> messageVOList.add(new MessageVO(message)));
        return messageVOList;
    }

    @PostMapping(value = "/savePdfInfo")
    public List<Document> savePdfInfo( MultipartFile  file) {
        //1.创建pdf读取器
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                file.getResource(),
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults()) //对pdf格式进行改造，去掉页眉之类的
                        .withPagesPerDocument(1) //对pdf进行分割，没1页作为一个片段
                        .build()
        );
        //2.读取pdf文档
        List<Document> documents = reader.read();
        //3.写入向量库
        vectorStore.add(documents);
        return vectorStore.similaritySearch("函数的命名规则");
    }

    @GetMapping(value = "argChat")
    public String argChat(@RequestParam("prompt") String prompt,@RequestParam("chatId") String chatId) {
        return pdfChatClient
                .prompt()
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,chatId))
                .advisors(advisorSpec -> advisorSpec.param(QuestionAnswerAdvisor.FILTER_EXPRESSION,"file_name == 'test.pdf'"))
                .call()
                .content();
    }

}
